from mpl_toolkits.mplot3d import axes3d
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
def plot(name):
    fig = plt.figure()
    ax = fig.gca(projection = '3d')
    data = np.loadtxt(name)
    x, y = np.meshgrid(np.arange(data.shape[1]), np.arange(data.shape[0]))
    surf = ax.plot_surface(x, y, data, rstride = 1, cstride = 1, cmap = cm.coolwarm, linewidth = 0, antialiased = False)
    ax.set_zlim(data.min(), data.max())
    fig.colorbar(surf, shrink = 0.5, aspect = 5)
    plt.savefig(name.split('.')[0] + ".png")
    plt.close()
